from com.bosssoft.handler.DataPreHandler import DataPreHandler
from com.bosssoft.handler.ModelTrainHandler import ModelTrainHandler
from com.bosssoft.utils.SimilarityUtil import SimilarityUtil
from com.bosssoft.utils.FileUtil import FileUtil
from datetime import datetime

if __name__ == '__main__':
    # 读入训练集---输入文本文件
    fileUtil = FileUtil(r"/Users/apple/PycharmProjects/sentence-similarity/data/test1/train.txt")
    trainSentences = fileUtil.read_lines()

    # 读入测试集----输入目标句子
    fileUtil = FileUtil(r"/Users/apple/PycharmProjects/sentence-similarity/data/test1/test.txt")
    testSentences = fileUtil.read_lines()

    # 训练模型-----通过调整参数:优化模型效率
    now = datetime.now()
    formatted_now = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"训练开始，{formatted_now}")
    trainHandler = ModelTrainHandler(DataPreHandler(trainSentences))
    modelInfo = trainHandler.TfidfModel() #28秒
    #modelInfo = trainHandler.LsiModel()  #40秒
    #modelInfo = trainHandler.LdaModel()  #48秒
    now = datetime.now()
    formatted_now = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"训练完成，开始测试，{formatted_now}")

    # 测试模型，将模型信息传入相似性测试工具
    similarityUtil = SimilarityUtil(modelInfo)
    # 与测试集中每个句子最相似的Top10个句子----输出与目标句子最相关的topN句子
    for i, test_word in enumerate(testSentences):
        top_n_sentences = similarityUtil.topNSimilarity(test_word, 2)
        print("=========与【" + test_word + "】相似的10个句子=========")
        for j, sentence in enumerate(top_n_sentences):
            print(f"{trainSentences[sentence.id]} => {sentence.score}")
        print("===================================================================================")


    now = datetime.now()
    formatted_now = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"测试结束，{formatted_now}")